import torch.nn as nn
from model.attack import Attack
from model.encoder import Encoder
from model.decoder import Decoder
from options import HiDDenConfiguration
from noise_layers.noiser import Noiser
from noise_layers.jpeg import Jpeg, JpegMask
from noise_layers.colorjitter import ColorJitter

class EncoderDecoder(nn.Module):
    """
    Combines Encoder->Noiser->Decoder into single pipeline.
    The input is the cover image and the watermark message. The module inserts the watermark into the image
    (obtaining encoded_image), then applies Noise layers (obtaining noised_image), then passes the noised_image
    to the Decoder which tries to recover the watermark (called decoded_message). The module outputs
    a three-tuple: (encoded_image, noised_image, decoded_message)
    """
    def __init__(self, config: HiDDenConfiguration):

        super(EncoderDecoder, self).__init__()
        self.encoder = Encoder(config)
        self.noiser = Attack(channels=64, attack=True)
        self.jp = Jpeg()
        # self.colorJitter = ColorJitter('Brightness')
        # self.jp = JpegMask(Q=50)
        self.decoder = Decoder(config)

    def forward(self, image, message):
        encoded_image = self.encoder(image, message)
        noised_image = self.noiser(encoded_image)
        # noised_image = self.jp(encoded_image)
        decoded_message = self.decoder(noised_image)
        return encoded_image, noised_image, decoded_message
